#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def load_prices(filename):
    """加载价格数据"""
    df = pd.read_csv(filename, sep=r'\s+', header=None, index_col=None)
    return df.values  # 转置，使行为时间，列为工具

def plot_all_stocks_with_stats(prices):
    """所有股票 + 统计信息"""
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(20, 16))
    
    # 上图：所有股票价格
    colors = plt.cm.tab20(np.linspace(0, 1, prices.shape[1]))
    for i in range(prices.shape[1]):
        ax1.plot(prices[:, i], 
                alpha=0.7, 
                linewidth=1,
                color=colors[i % len(colors)])
    
    ax1.set_title('50支股票价格走势', fontsize=14, fontweight='bold')
    ax1.set_xlabel('时间 (交易日)')
    ax1.set_ylabel('价格')
    ax1.grid(True, alpha=0.3)
    
    # 下图：统计信息
    mean_prices = np.mean(prices, axis=0)
    std_prices = np.std(prices, axis=0)
    
    x_pos = np.arange(prices.shape[1])
    bars = ax2.bar(x_pos, mean_prices, alpha=0.7, color='skyblue')
    
    # 添加误差线（标准差）
    ax2.errorbar(x_pos, mean_prices, yerr=std_prices, fmt='none', color='red', alpha=0.5)
    
    ax2.set_title('各股票平均价格 & 标准差', fontsize=14, fontweight='bold')
    ax2.set_xlabel('股票编号')
    ax2.set_ylabel('价格')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    ax2.plot(x_pos, mean_prices, color='black', linewidth=2, label='均值线')
    ax2.legend()
    plt.show()

def main():
    """主函数"""
    # 加载数据
    print("正在加载价格数据...")
    prices = load_prices('prices.txt')
    print(f"成功加载数据: {prices.shape[0]} 个时间点, {prices.shape[1]} 支股票")
    
    print("生成带统计信息版全股票图...")
    plot_all_stocks_with_stats(prices)

if __name__ == "__main__":
    main()